import os
import numpy as np

from tqdm import tqdm
from typing import Any, Dict, List
from .structure import RadarPoint
from .i_reader import IReader


class RadarReader(IReader[RadarPoint]):
    def __init__(self, cfg: Dict[str, Any]):
        self.point_size = cfg.get('point_size')
        self.parallel = cfg.get('parallel')
        self.data_batch_size = cfg.get('data_batch_size')

        self.z_min, self.z_max = cfg.get('height_range')

    def _get_file_list(self, folder_path) -> List[str]:
        files = []
        with os.scandir(folder_path) as it:
            for entry in it:
                if entry.name.endswith('.bin') and entry.is_file():
                    path = os.path.join(folder_path, entry.name)
                    file_size = entry.stat().st_size
                    if file_size % self.point_size != 0:
                        raise ValueError(f"Invalid file size: {path} ({file_size} bytes)")
                    files.append(path)
        return files

    def _read_single_file(self, file_path: str) -> List[RadarPoint]:
        with open(file_path, 'rb') as f:
            buffer = f.read()
        
        data = np.frombuffer(buffer, dtype=np.float32).reshape(-1, self.point_size // 4)

        mask = (
            (data[:, 2] >= self.z_min) &
            (data[:, 2] <= self.z_max)
        )
        data = data[mask]

        return [
            RadarPoint(
                strength=point[0],
                x=point[1],
                y=point[3], 
                z=point[2],
                velocity=point[4],
                cohen_factor=point[5]
            ) for point in data
        ]

    def process_all_files(self, folder_path, batch) -> Dict[str, List[RadarPoint]]:
        # 获取文件夹中的文件列表
        files = self._get_file_list(folder_path)
        # 按文件名排序
        files.sort(key=lambda x: int(x.split('_')[-1].split('.')[0]))
        files = files[(batch-1)*self.data_batch_size:batch*self.data_batch_size] 
        
        if self.parallel:
            from multiprocessing import Pool, cpu_count

            with Pool(max(1, cpu_count()-1)) as pool:
                results = list(tqdm(
                    pool.imap(self._read_single_file, files),
                    total=len(files),
                    desc="Reading bin files"
                ))
        else:
            results = []
            for f in tqdm(files, desc="Reading bin files"):
                try:
                    results.append(self._read_single_file(f))
                except ValueError as e:
                    print(f"Skipping file {f}: {str(e)}")
                    continue
        
        points = {os.path.splitext(os.path.basename(f))[0]: res for f, res in zip(files, results)}
        return points
